import random
import numpy as np

from beam.util_psll import get_psll_by_phase


class RISHarrisHawksOptimization:
    __bit_num = 0  # 比特数
    __beam_num = 0  # 波束数

    __nRow = 32  # x方向单元个数，θ方向
    __mLine = 32  # y方向单元个数，φ方向

    def __init__(self, bit_num, beam_num, num_hawks=50, max_iter=100, lb=-180, ub=180, phase_mix_init=None):
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        self.num_hawks = num_hawks  # 鹰的数量
        self.max_iter = max_iter  # 最大迭代次数
        if phase_mix_init is not None:
            self.lb = np.full(phase_mix_init.shape, lb)  # 搜索空间下界
            self.ub = np.full(phase_mix_init.shape, ub)  # 搜索空间上界
        else:
            self.lb = np.full((self.__nRow, self.__nRow), lb)  # 搜索空间下界
            self.ub = np.full((self.__nRow, self.__nRow), ub)  # 搜索空间上界
        self.hawks = None  # 鹰的位置
        self.best_fitness = float('inf')  # 最佳适应度
        self.best_position = None  # 最佳位置
        self.iteration = 0  # 当前迭代次数
        self.best_fitness_history = []      # 保存每一代的最佳适应度值
        self.best_individual_history = []   # 保存每一代的最佳适应度个体

    def fitness(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def initialize_population(self, phase_mix_init):
        # 初始化鹰的位置
        self.hawks = np.random.uniform(self.lb, self.ub, (self.num_hawks, *phase_mix_init.shape))
        # 将初始位置设置为第一个鹰的位置
        self.hawks[0] = phase_mix_init
        # 确保初始位置在边界内
        self.hawks[0] = np.clip(self.hawks[0], self.lb, self.ub)
        # 计算初始适应度
        fitness_values = [self.fitness(hawk) for hawk in self.hawks]
        best_index = np.argmin(fitness_values)
        self.best_fitness = fitness_values[best_index]
        self.best_position = self.hawks[best_index].copy()

    def update_position(self):
        # 更新位置
        for i in range(self.num_hawks):
            E0 = 2 * (1 - (self.iteration / self.max_iter))
            E = 2 * E0 * np.random.rand() - E0
            J = 2 * (1 - np.random.rand())

            if abs(E) >= 1:
                # 软包围
                X_rand = self.hawks[np.random.randint(0, self.num_hawks)]
                D_X_rand = abs(J * X_rand - self.hawks[i])
                self.hawks[i] = X_rand - D_X_rand
            else:
                # 硬包围
                S = self.best_position - self.hawks[i]
                if np.random.rand() < 0.5:
                    # 攻击模式1
                    D_S = abs(S)
                    L = np.random.uniform(-1, 1, size=self.best_position.shape)
                    self.hawks[i] = self.best_position - E * D_S * L
                else:
                    # 攻击模式2
                    Y = self.hawks[np.random.randint(0, self.num_hawks)]
                    Z = self.hawks[np.random.randint(0, self.num_hawks)]
                    Y_new = Y - (np.random.rand() * abs(Y - Z))
                    S_Y = self.best_position - Y_new
                    self.hawks[i] = Y_new - E * abs(S_Y)

            # 应用边界条件
            self.hawks[i] = np.clip(self.hawks[i], self.lb, self.ub)

    def run(self, phase_mix_init, logger):
        logger.info("num_hawks=%d, max_iter=%d" % (self.num_hawks, self.max_iter))

        # 如果在初始化时没有提供 phase_mix_init，则在这里重新初始化
        if self.lb is None or self.ub is None:
            self.lb = np.full(phase_mix_init.shape, -180)  # 搜索空间下界
            self.ub = np.full(phase_mix_init.shape, 180)  # 搜索空间上界

        self.best_individual = phase_mix_init
        self.best_fitness = self.fitness(phase_mix_init)
        self.best_fitness_history = []
        self.best_individual_history = []

        # 初始化返回值
        self.initialize_population(phase_mix_init)

        for self.iteration in range(self.max_iter):
            self.update_position()
            # 计算当前适应度
            fitness_values = [self.fitness(hawk) for hawk in self.hawks]
            best_index = np.argmin(fitness_values)
            current_best_fitness = fitness_values[best_index]

            if current_best_fitness < self.best_fitness:
                self.best_fitness = current_best_fitness
                self.best_position = self.hawks[best_index].copy()

            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_position.copy())

            logger.info("iteration=%d: self.best_fitness=%f" % (self.iteration, self.best_fitness))

        # 返回最佳个体、最佳适应度、历史最佳适应度和历史最佳个体
        return self.best_position, self.best_fitness, self.best_fitness_history, self.best_individual_history
